import networkx as nx
from timer import Timerx
from itertools import izip
import numpy as np
from sklearn.utils import shuffle

def read_graph(file_path, delimiter=' ', convert_string_to_int = True):
    G = nx.Graph()
    
    with open(file_path) as edges_file:
        for line in edges_file:
            source, target = line.rstrip().split(delimiter)
            if (convert_string_to_int):
                G.add_edge(int(source), int(target))
            else:
                G.add_edge(source, target)
    
    return G

def build_U(nodes_list):
    """
    Returns a list U which is the universal edges set that contains 
    every possible edge that can be in the graph. 
    The size of this set N*(N-1)/2 where N is the number of nodes/vertices.
    """
    
    N = len(nodes_list)
    U_size = (N*(N-1))/2
    U = np.zeros((U_size,2))
    
    j = 0
    for i, node1 in enumerate(nodes_list):
        for node2 in nodes_list[i+1:]:
            edge = (node1,node2)
            U[j] = edge
            j += 1
            
    #Old fast but memory consuming way
#     U = list()
#  
#     for i, node1 in enumerate(nodes_list):
#         for node2 in nodes_list[i+1:]:
#             edge = (node1,node2)
#             U.append(edge)
    
    return U

def build_small_U(G):
    """
    Returns a down sampled list U, where U is a list that contains tuples of edges.
    For example U[0] = (node_i, node_j). Down sampled means here is that U contains 
    the same number of (edges that truly exists) AND (edges that don't exist) ... (i.e. positive
    and negative examples for training).
    
    Parameters:
    -----------
    G: a networkx graph object.
    
    Returns:
    --------
    U: a list of edges.
    """
    U = G.edges()
    sample_size = len(U)
    
    A = nx.adj_matrix(G)
    A = np.array(A)
    
    nonedges = np.transpose(np.nonzero(A == 0))
    nonedges = shuffle(nonedges, random_state=0)
    
    index_to_node = {}
    for i, node in enumerate(G.nodes()):
        index_to_node[i] = node
    
    for x, (i, j) in enumerate(nonedges):
        if x < sample_size:
            node1 = index_to_node[i]
            node2 = index_to_node[j]
            U.append( (node1, node2) )
        else:
            break
        
    return U
        
    

def build_Y(G, U):
    """
    Builds the Y ground-truth list that contains a label
    for each edge in the universal to indicate if the edge
    really exists in the graph or not.
    
    Parameters:
    G: the Networkx graph.
    U: the universal edges set that contains every possible edge in the graph.
    """
    N = len(U)
    Y = np.zeros(N)
    
    for i, edge in enumerate(U):
        node1 = edge[0]
        node2 = edge[1]
        if G.has_edge(node1,node2):
            Y[i] = 1
            
#     Y = list()
#     
#     for edge in U:
#         node1 = edge[0]
#         node2 = edge[1]
#         if G.has_edge(node1,node2):
#             Y.append(1)
#         else:
#             Y.append(0)
    
    return Y

def add_feature(data, U, S, nodes_list):
    """
    Extends the data matrix with a column that represents 
    a new feature.
    
    Parameters:
    data: the data matrix (usually called X in sci-kit).
    U: the universal edges list (set).
    S: the similarity matrix that holds the similarities between each
       node in the graph.
    nodes_list: the list of nodes from a networkx Graph (returned by the G.nodes() method). 
       
    Returns the data extended with a new column.
    """
    
    nodes = {}
    for i, node in enumerate(nodes_list):
        nodes[node] = i
    
    N = len(U)
    feature = np.zeros(N)
    
    for i, edge in enumerate(U):
        node1 = edge[0]
        node2 = edge[1]
        node1_index = nodes[node1]
        node2_index = nodes[node2]
        feature[i] = S[node1_index, node2_index]

    if data == None:
        data = np.reshape(feature, (-1,1))
    else:
        data = np.column_stack([data, feature])

    return data

def add_raw_feature(data, U, M, nodes_list):
    """
    Extends the data matrix (X in scikit) with raw features from M.
    So each dyad (i,j) features is just the concatination of the row
    of i in M and the row of j in M.
    
    Parameters
    ----------
    data: features numpy matrix.
    U: the universal edges set.
    M: the raw features set (this is most likely to be the adjacency matrix)
    nodes_list: a list of all nodes in the graph

    Returns
    -------
    data: the features numpy matrix extended with columns
    """
    nodes = {}
    for i, node in enumerate(nodes_list):
        nodes[node] = i
    
    N = len(U)
    C = M.shape[1] #num of columns
    feature = np.zeros((N, C*2)) #number dyads times the number of double coulmns
    
    for i, edge in enumerate(U):
        node1 = edge[0]
        node2 = edge[1]
        node1_index = nodes[node1]
        node2_index = nodes[node2]
        feature[i] = np.concatenate(( M[node1_index], M[node2_index] ))
        
    if data == None:
        data = feature
    else:
        data = np.column_stack([data, feature])
        
    return data

def node_neighborhood(G, node, n):
    """
    Returns a list of nodes which are the n-neighborhood of the input node.
    
    Parameters
    ----------
    G: networkx graph object.
    node: the node to get the neighborhood for.
    n: the neighborhood degree.
    """
    path_lengths = nx.single_source_dijkstra_path_length(G, node)
    return [node for node, length in path_lengths.iteritems()
                    if length == n]

def build_Un_neighborhoods(G, U, lengths_list):
    """
    Builds U_n the set of all possible links for nodes that are within 
    a specified neighborhood length n.
    
    Parameters
    ----------
    G: networkx graph object.
    U: the set of all possible links.
    lengths_list: a list that holds the neighborhoods degrees to compute.
    
    Returns
    -------
    a dictionary mapped by neighborhood degree. The values are numpy 2D arrays
     where each vector is of the form [v_i, v_j]
    to indicate that v_i has a shortest path of length n to v_j
    """
    
    U_n = {}
    for n in lengths_list:
        U_n[n] = []
        
    paths_lengths = nx.shortest_path_length(G)
    
    for edge in U:
        node1 = edge[0]
        node2 = edge[1]
        if paths_lengths[node1].has_key(node2):
            plength = paths_lengths[node1][node2]
        else:
            plength = 0
        if  U_n.has_key(plength):
            U_n[plength].append(edge)
    
    for k, v in U_n.iteritems():
        U_n[k] = np.array(v)
    
    return U_n
    
    
class DyadFeaturesIndexLocator:
    """
    Returns the index of a given dyad in the 
    dyads features matrix.
    """
    
    def __init__(self, U, nodes_list):
        """
        U: the set of all possible edges/dyads
        node_list: the list of all nodes in the graph
        """
        
        self.nodes = {}
        for i, node in enumerate(nodes_list):
            self.nodes[node] = i
        
        self.M = {}
        for i, edge in enumerate(U):
            node1 = edge[0]
            node2 = edge[1]
            node1_index = self.nodes[node1]
            node2_index = self.nodes[node2]
            self.M[ (node1_index, node2_index) ] = i
    
    def get_feature_index(self, node1_index, node2_index):
        """
        Returns the index of a given dyad in the
        dyads features matrix.
        
        Parameters:
        -----------
        node1_index: the index of the first node 
        node2_index: the index of the secod node
        """
        if (node1_index, node2_index) in self.M:
            index = self.M[ (node1_index, node2_index) ]
        elif (node2_index, node1_index) in self.M:
            index = self.M[ (node2_index, node1_index) ]
        else:
            raise Exception("I can't find an index for the given\
             dyad in the dyad feature matrix!!!")
        return index
    
    
def ttt():
    U = [(1,2), (1,3), (2,3)]
    nodes_list = [1,2,3]
    X = [12,13,23]
    locator = DyadFeaturesIndexLocator(U, nodes_list)
    index =  locator.get_feature_index(0,1)
    print X[index]
    
def get_extended_graph(G, X, bins=None):
    """
    Extends the networkx graph G with nodes that represent attributes and 
    links between the attribute nodes and normal nodes if the normal nodes have
    the attribute values. A Real-valued attribute is divided into bins where each bin is a node, and 
    a link between a normal node and a bin node is created if the normal node has an attribute value 
    that falls within the bin node range.

    Similarly each binary attribute is represented by two nodes, a node that represents a value 1 and 
    another that represents a value 0.
    
    Parameters:
    -----------
    G: a networkx graph object.
    X: a design matrix where columns are attributes and each row represents a node.
       So each row is the feature vector of a node.
       
    Returns:
    ----------
    extended_G: returns an extended version of G with attributes nodes and the 
               links between them and normal nodes.

    """
    G = G.copy()
    if bins == None:
        bins = [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
        bins = np.array(bins)
    
    index_to_node = {} #to map from index to node name
    for i, node in enumerate(G.nodes()):
        index_to_node[i] = node

    n_rows = X.shape[0]
    n_cols = X.shape[1]

    
    for col in xrange(n_cols):
        #check if X[:, col] is binary or real
        v = (X[:,col] == 1)

        num_ones = np.where(v)[0].size
        
        v = (X[:,col] == 0)
        
        num_zeros = np.where(v)[0].size
        
        total_size = num_ones + num_zeros

        is_binary = False
        
        if n_rows == total_size:
            is_binary = True
            
        if is_binary:
            node0 = "col_" + str(col) + "_0" 
            node1 = "col_" + str(col) + "_1" 
            
            zeroval_nodes_inds =  np.where((X[:,col] == 0))[0] #returns the indices of the nodes that has the attribute value as 0
            onesval_nodes_inds =  np.where((X[:,col] == 1))[0] #returns the indices of the nodes that has the attribute value as 1

            nodes_linked_to_zero = [ index_to_node[node_index] for node_index in  zeroval_nodes_inds ]
            nodes_linked_to_one = [ index_to_node[node_index] for node_index in  onesval_nodes_inds ]
            
            zeros_edges = [ (node0, x) for x in nodes_linked_to_zero ]
            ones_edges = [ (node1, x) for x in nodes_linked_to_one ]

            if len(zeros_edges) > 0:
                G.add_edges_from(zeros_edges)
                
            if len(ones_edges) > 0:
                G.add_edges_from(ones_edges)
        else:
            column = X[:,col]
            binplace = np.digitize(column, bins)#each element in binplace is the index of the bin that this element falls into.
            for bin in xrange(bins.size):
                elements_inds = np.where(binplace==bin)[0]
                nodes_linked_to_bin = [ index_to_node[node_index] for node_index in  elements_inds ]
                node_bin = "col_" + str(col) + "_" + str(bin)
                edges =  [ (node_bin, x) for x in nodes_linked_to_bin ]
                if len(edges) > 0:
                    G.add_edges_from(edges)



    return G

